package main

import (
	"flag"
	"fmt"
	"log"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/zeromicro/go-zero/core/limit"
	"github.com/zeromicro/go-zero/core/stores/redis"
)

var (
	rdx     = flag.String("redis", "192.168.8.20:6379", "the redis, default localhost:6379")
	rdxKey  = flag.String("redisKey", "rate", "the redis key, default rate")
	rdxPass = flag.String("redisPass", "admin", "the redis password")
	threads = flag.Int("threads", runtime.NumCPU(), "the concurrent threads, default to cores")
)

func main() {
	flag.Parse()
	const seconds = 1
	store := redis.MustNewRedis(redis.RedisConf{Host: *rdx, Type: redis.NodeType}, redis.WithPass(*rdxPass))
	fmt.Println(store.Ping())
	//每s生产 rate 消耗burst
	lmt := limit.NewTokenLimiter(1, 10, store, *rdxKey)
	//lmt = limit.NewTokenLimiter(10, 1, store, *rdxKey)
	timer := time.NewTimer(time.Second * seconds)
	quit := make(chan struct{})
	defer timer.Stop()
	go func() {
		<-timer.C
		close(quit)
	}()

	var allowed, denied int32
	var wait sync.WaitGroup
	log.Println("threads", *threads)
	for i := 0; i < *threads; i++ {
		log.Println(i)
		wait.Add(1)
		go func() {
			for {
				select {
				case <-quit:
					wait.Done()
					return
				default:
					if lmt.Allow() {
						atomic.AddInt32(&allowed, 1)
					} else {
						atomic.AddInt32(&denied, 1)
					}
				}
			}
		}()
	}

	wait.Wait()
	fmt.Printf("allowed: %d, denied: %d, qps: %d\n", allowed, denied, (allowed+denied)/seconds)
}
